#阿基米德优化算法
import numpy as np
import random
import math
import matplotlib.pyplot as plt
import time
import initialization_map

class AOA:
    def __init__(self,n_dim=2, pop_size=10, max_iter=100, lb=[0], ub=[50],obj_func=None,C3=1,C4=1,ax=None):
        self.n_dim = n_dim
        self.pop_size = pop_size
        self.max_iter = max_iter
        self.lb, self.ub = lb, ub
        self.func = obj_func
        self.C1=2
        self.C2=6
        self.C3 = C3
        self.C4 = C4
        self.u=.9
        self.l=.1   #paramters in Eq. (12)
        self.acc_temp = []
        self.ax = ax
        self.X = np.zeros([pop_size, n_dim])
        self.acc = np.zeros([pop_size, n_dim])
        for i in range(pop_size):
            for j in range(n_dim):
                self.X[i, j] = np.random.random() * (ub[j] - lb[j]) + lb[j]
                self.acc[i,j] = np.random.random() * (ub[j] - lb[j]) + lb[j]
        self.den=np.random.rand(pop_size,n_dim)
        self.vol=np.random.rand(pop_size,n_dim)
        # self.acc=self.lb+np.random.rand(pop_size,n_dim)*(self.ub-self.lb)
        self.Y = np.zeros([pop_size,1])
        for i in range(pop_size):
            self.Y[i]=obj_func(self.X[i])
        self.Scorebest = np.min(self.Y)
        self.Score_index = np.argmin(self.Y)
        self.Xbest = self.X[self.Score_index].copy()
        self.den_best=self.den[self.Score_index].copy()
        self.vol_best=self.vol[self.Score_index].copy()
        self.acc_best=self.acc[self.Score_index].copy()
        self.acc_norm=self.acc.copy()
        self.Convergence_curve = np.zeros(max_iter+1)
        self.Convergence_curve[0] = self.Scorebest
        if (self.ax):
            # 初始化标点
            self.ax.chat3dplot(self.X, self.func, 'ax1', 'o', 'black')

    def update(self):
        Max_iter = self.max_iter
        Materials_no = self.pop_size
        den = self.den
        vol = self.vol
        den_best = self.den_best
        vol_best = self.vol_best
        acc = self.acc
        acc_best = self.acc_best
        avg_fit = []
        avg_ts = []#平均方差
        time_start = time.time() #记录迭代寻优开始时间
        for t in range(Max_iter):
            TF=np.exp(((t-Max_iter)/(Max_iter)))#迁移算子，区分全局探索和局部开发
            if TF>1:
                TF=1
            d=np.exp((Max_iter-t)/Max_iter)-(t/Max_iter)#密度因子
            self.acc = self.acc_norm.copy()
            r=np.random.rand()
            #先更新加速度
            for i in range(Materials_no):
                #每代个体密度和体积的更新
                den[i]=den[i]+r*(den_best-den[i,:])
                vol[i]=vol[i]+r*(vol_best-vol[i,:])
                if TF<=.5:#全局搜索
                    mr=np.random.choice(Materials_no)
                    if len(self.acc_temp) < Materials_no:
                        self.acc_temp.append((den[mr] + (vol[mr] * acc[mr])) / (den[i] * vol[i]) ) # Eq. (10)
                    else:
                        self.acc_temp[i] = (den[mr] + (vol[mr] * acc[mr])) / (den[i] * vol[i]) # Eq. (10)
                else:#局部搜索
                    if len(self.acc_temp) < Materials_no:
                        self.acc_temp.append((den_best + (vol_best * acc_best)) / (den[i] * vol[i]) )  # Eq. (11)              
                    else:
                        self.acc_temp[i] = (den_best + (vol_best * acc_best)) / (den[i] * vol[i]) # Eq. (11)
            #为了规范个体更新步长范围，加速度归一化操作
            self.acc_norm=((self.u*(self.acc_temp-np.min(self.acc_temp)))/(np.max(self.acc_temp)-np.min(self.acc_temp)))+self.l;   # Eq. (12)
            Xnew = np.zeros((self.X.shape))
            #再用加速度更新x位置
            for i in range(Materials_no):
                if TF<.5:
                    for j in range(self.X.shape[1]):
        
                        mrand=np.random.choice(Materials_no)
                        Xnew[i][j]=self.X[i][j]+self.C1*np.random.rand()*self.acc_norm[i][j]*(self.X[mrand][j]-self.X[i][j])*d  # Eq. (13)
                else:
                    for j in range(self.X.shape[1]):
                        p=2*np.random.rand()-self.C4  # Eq. (15)
                        T=self.C3*TF
                        if T>1:
                            T=1
                        if p<.5:
                            Xnew[i][j]=self.Xbest[j]+self.C2*np.random.rand()*self.acc_norm[i][j]*(T*self.Xbest[j]-self.X[i][j])*d  # Eq. (14)
                        else:
                            Xnew[i][j]=self.Xbest[j]-self.C2*np.random.rand()*self.acc_norm[i][j]*(T*self.Xbest[j]-self.X[i][j])*d
            Xnew = self.fun_checkpositions(self.n_dim,Xnew,Materials_no,self.lb,self.ub).copy()
            for i in range(Materials_no):
                v = self.func(Xnew[i])
                if v < self.Y[i]:
                    self.X[i] = Xnew[i]
                    self.Y[i] = v
            var_Ybest = np.min(self.Y)
            var_index = np.argmin(self.Y)
            self.Convergence_curve[t+1] = var_Ybest
            # 展示描点图
            if self.ax and t == self.ax.a1:
                # 最终结果化标点
                self.ax.chat3dplot(self.X, self.func, 'ax2', 'o', 'black')

            elif self.ax and t == self.ax.a2:
                self.ax.chat3dplot(self.X, self.func, 'ax3', 'o', 'black')
            if var_Ybest < self.Scorebest:
                self.Scorebest = var_Ybest
                self.Score_index = var_index
                self.Xbest = self.X[var_index].copy()
                self.den_best = self.den[self.Score_index].copy()
                self.vol_best = self.vol[self.Score_index].copy()
                self.acc_best = self.acc_norm[self.Score_index].copy()
            
            avg_fit.append(np.mean(self.Y))
            avg_ts.append(np.var(self.Y))
        time_end = time.time() #记录迭代结束时间
        print(f'AOA共花费 {time_end - time_start} 秒')
        if self.ax:
            # 最终结果化标点
            self.ax.chat3dplot(self.X, self.func,'ax4', 'o', 'black')
            plt.show()
        print('AOA最优适应度', self.Scorebest)
        print('AOA最优解', self.Xbest)
        return avg_fit,avg_ts
        
    def fun_checkpositions(self,dim,vec_pos,var_no_group,lb,ub):
        for i in range(var_no_group):
            for j in range(dim):
                if vec_pos[i][j] < lb[j]:
                    vec_pos[i][j] = lb[j]
                if vec_pos[i][j] > ub[j]:
                    vec_pos[i][j] = ub[j]
    
        return vec_pos
    def aoa(self,isBox = False):
        avg_fit,avg_ts = self.update()
        if isBox:
            return self.Convergence_curve, self.X
        else:
            return self.Convergence_curve, avg_fit, avg_ts



#基准函数测试
# def sphereModel(variables_values = [0,0]):
#     func_value = 0
    
#     for i in range(len(variables_values)):
#         func_value += variables_values[i]**2
#     return func_value



# if __name__ == '__main__':
#     low = [-100, -100]
#     up = [100,100]
#     aoa = AOA(2,10,100,low,up,sphereModel)
#     Convergence_curve = aoa.aoa()
#     print(f'测试{Convergence_curve}')
#     plt.semilogy(Convergence_curve)
#     plt.show()